package com.example.auth.config;

import com.example.auth.util.JwtUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import java.security.Principal;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    private static final Logger logger = LoggerFactory.getLogger(WebSocketConfig.class);

    private final JwtUtil jwtUtil;

    public WebSocketConfig(JwtUtil jwtUtil) {
        this.jwtUtil = jwtUtil;
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        // Register STOMP endpoints, add SockJS support
        logger.info("Registering WebSocket endpoint: /ws");
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")  // Allow all origins
                .setAllowedOrigins("http://localhost:5173")  // Explicitly add Vite development server
                .withSockJS()
                .setClientLibraryUrl("https://cdn.jsdelivr.net/npm/sockjs-client@1/dist/sockjs.min.js")  // Use SockJS client library from CDN
                .setWebSocketEnabled(true)  // Ensure WebSocket transport is enabled
                .setSessionCookieNeeded(false);  // Disable session cookie, simplifying development
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        // Configure message broker
        logger.info("Configuring WebSocket message broker");

        // Destination prefix for client-sent messages
        registry.setApplicationDestinationPrefixes("/app");
        // Destination prefix for server-sent messages to clients
        registry.enableSimpleBroker("/user", "/topic");
        // Set prefix for point-to-point communication
        registry.setUserDestinationPrefix("/user");

        logger.info("WebSocket message broker configuration completed");
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        logger.info("Configuring WebSocket client inbound channel interceptor");

        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

                if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {
                    logger.info("Received STOMP CONNECT command");
                    String token = accessor.getFirstNativeHeader("Authorization");

                    if (token != null && token.startsWith("Bearer ")) {
                        token = token.substring(7);
                        logger.info("Extracted JWT Token from WebSocket connection: {}", token.substring(0, Math.min(10, token.length())) + "...");

                        // Validate Token
                        if (jwtUtil.validateToken(token)) {
                            try {
                                // Extract user ID from token
                                Long userId = jwtUtil.getUserIdFromToken(token);
                                logger.info("User ID extracted from JWT Token: {}", userId);

                                if (userId != null) {
                                    // Set user ID as Principal name
                                    final Long finalUserId = userId;
                                    accessor.setUser(() -> String.valueOf(finalUserId));
                                    logger.info("Successfully set user ID as Principal for WebSocket connection: {}", userId);
                                } else {
                                    logger.error("WebSocket connection failed: Token does not contain a valid user ID");
                                    return null; // Reject connection
                                }
                            } catch (Exception e) {
                                logger.error("Failed to set Principal for WebSocket connection: {}", e.getMessage(), e);
                                return null; // Reject connection
                            }
                        } else {
                            logger.error("WebSocket connection JWT validation failed, Token is invalid");
                            return null; // Reject connection
                        }
                    } else {
                        logger.error("WebSocket connection missing Authorization header or incorrect format");
                        return null; // Reject connection
                    }
                } else if (accessor != null && StompCommand.DISCONNECT.equals(accessor.getCommand())) {
                    Principal user = accessor.getUser();
                    if (user == null) {
                        logger.error("Missing user identity in message, rejecting message: {}", accessor.getCommand());
                        return null;
                    }
                }
                return message;
            }
        });

        logger.info("WebSocket client inbound channel interceptor configuration completed");
    }
}